﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using OpenCvSharp;
using OpenCvSharp.Dnn;
using OpenCvSharp.Extensions;

namespace VisualTargetDetection
{
    /// <summary>
    /// 检测结果
    /// </summary>
    public struct DetectionResult
    {
        /// <summary>
        /// 标注编号
        /// </summary>
        public int Lable;
        /// <summary>
        /// 标注名称
        /// </summary>
        public string Name;
        /// <summary>
        /// 包含位置和大小信息的长方形
        /// </summary>
        public Rect2d Rectangle;
        /// <summary>
        /// 检测项目的相似程度和识别把握
        /// </summary>
        public float Probability, Confidence;
    }

    public class Detector
    {
        private string Cfg, Weight, Names;
        private string[] Labels;
        private Scalar[] Colors = Enumerable.Repeat(false, 2).Select(x => Scalar.RandomColor()).ToArray();
        private Net net;

        /// <summary>
        /// 检测器类
        /// <para>加载模型，输入一帧Mat画面，输出目标信息，包含位置、大小和Lable</para>
        /// </summary>
        /// <param name="config">.cfg配置文件位置</param>
        /// <param name="weight">.weights权重文件位置</param>
        /// <param name="names">.txt标注名称文件位置</param>
        public Detector(string config = "./mod/yolov3.cfg", string weight = "./mod/yolov3.weights", string names = "./mod/classes.txt")
        {
            Cfg = Path.GetFullPath(config);
            Weight = Path.GetFullPath(weight);
            Names = Path.GetFullPath(names);
            Labels = File.ReadAllLines(names).ToArray();
            net = CvDnn.ReadNetFromDarknet(Cfg, Weight);
        }

        /// <summary>
        /// 进行目标检测，返回检测结果结构体
        /// </summary>
        /// <param name="imgSrc">一帧画面</param>
        public DetectionResult[] Detect(Mat imgSrc)
        {
            Mat org = new Mat();
            org = imgSrc;

            const float threshold = 0.5f;       //confidence阈值 
            const float nmsThreshold = 0.3f;    //nms阈值
            var blob = CvDnn.BlobFromImage(org, 1.0 / 255, new OpenCvSharp.Size(416, 416), new Scalar(), true, false);
            net.SetInput(blob);
            var outNames = net.GetUnconnectedOutLayersNames();
            var outs = outNames.Select(_ => new Mat());
            net.Forward(outs, outNames);
            return _Detect(outs, org, threshold, nmsThreshold);
        }


        private DetectionResult[] _Detect(IEnumerable<Mat> output, Mat image, float threshold, float nmsThreshold, bool nms = true)
        {
            List<DetectionResult> resultlist = new();

            //NMS变量
            var classIds = new List<int>();
            var confidences = new List<float>();
            var probabilities = new List<float>();
            var boxes = new List<Rect2d>();

            var w = image.Width;
            var h = image.Height;
            /*
             YOLO3 COCO trainval output
             0 1 : center                    2 3 : w/h
             4 : confidence                  5+: class probability 
            */
            const int prefix = 5;//class probability

            foreach (var prob in output)
            {
                for (var i = 0; i < prob.Rows; i++)
                {
                    var confidence = prob.At<float>(i, 4);
                    if (confidence > threshold) //阈值
                    {
                        double maxVal, minVal;
                        OpenCvSharp.Point min, max;
                        Cv2.MinMaxLoc(prob.Row(i).ColRange(prefix, prob.Cols), out minVal, out maxVal, out min, out max);
                        var classes = max.X;
                        var probability = prob.At<float>(i, classes + prefix);

                        if (probability > threshold) //阈值
                        {
                            //中心点坐标&大小
                            var centerX = prob.At<float>(i, 0) * w;
                            var centerY = prob.At<float>(i, 1) * h;
                            var width = prob.At<float>(i, 2) * w;
                            var height = prob.At<float>(i, 3) * h;

                            if (!nms)
                            {//关闭NMS
                                resultlist.Add(new DetectionResult
                                {
                                    Lable = classes,
                                    Rectangle = new Rect2d(centerX, centerY, width, height),
                                    Name = Labels[classes],
                                    Confidence = confidence,
                                    Probability = probability
                                });
                                continue;
                            }

                            //准备NMS数据
                            classIds.Add(classes);
                            confidences.Add(confidence);
                            probabilities.Add(probability);
                            boxes.Add(new Rect2d(centerX, centerY, width, height));
                        }
                    }
                }
            }

            if (!nms) return resultlist.ToArray();

            //NMS非极大值抑制，减少重复检出
            int[] indices;
            CvDnn.NMSBoxes(boxes, confidences, threshold, nmsThreshold, out indices);


            foreach (var i in indices)
            {
                var box = boxes[i];
                resultlist.Add(new DetectionResult
                {
                    Lable = classIds[i],
                    Rectangle = new Rect2d(box.X, box.Y, box.Width, box.Height),
                    Name = Labels[classIds[i]],
                    Confidence = confidences[i],
                    Probability = probabilities[i]
                });
            }
            return resultlist.ToArray();
        }
    }
}
